//
//  Trie.swift
//  Chapter3LinkedList
//
//  Created by TRS-Mobile on 2021/9/2.
//

import Foundation

class Trie<CollectionType: Collection> where CollectionType.Element: Hashable {
    typealias Node = TriedNode<CollectionType.Element>
    private let root = Node(key: .none, parent: .none)
    init() {}
    
    func insert(_ collection: CollectionType) -> Void {
        var current = root
        for item in collection {
            if current.children[item] == nil {
                current.children[item] = Node(key: item, parent: current)
            }
            current = current.children[item]!
        }
        current.isTerminating = true
    }
    
    func remove(_ collection: CollectionType) -> Void {
        var current = root
        for item in collection {
            guard let child = current.children[item] else { return }
            current = child
        }
        guard current.isTerminating else {
            return
        }
        current.isTerminating = false
        while let parent = current.parent, current.children.isEmpty && !current.isTerminating {
            parent.children[current.key!] = nil
            current = parent
        }
    }
    
    func contains(_ collection: CollectionType) -> Bool {
        var current = root
        for item in collection {
            guard let child = current.children[item] else { return false }
            current = child
        }
        return current.isTerminating
    }
}

extension Trie where CollectionType: RangeReplaceableCollection {
    
    private func collections(startingWith prefix: CollectionType,
                             after node: Node) -> [CollectionType] {
        var results: [CollectionType] = []
        if node.isTerminating {
            results.append(prefix)
        }
        for child in node.children.values {
            var prefix = prefix
            prefix.append(child.key!)
            results.append(contentsOf: collections(startingWith: prefix, after: child))
            
        }
        return results
    }
    
    func collections(startingWith prefix: CollectionType) -> [CollectionType] {
        var current = root
        for item in prefix {
            guard let child = current.children[item] else { return [] }
            current = child
        }
        return collections(startingWith: prefix, after: current)
    }
}
